import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class Ant:
    def __init__(self, num_nodes):
        self.num_nodes = num_nodes
        self.reset()

    def reset(self):
        self.path = []
        self.total_cost = 0.0
        self.visited = [False] * self.num_nodes

    def visit_node(self, node, cost_matrix):
        if len(self.path) > 0:
            self.total_cost += cost_matrix[self.path[-1], node]
        self.path.append(node)
        self.visited[node] = True

    def path_cost(self, cost_matrix):
        return self.total_cost + cost_matrix[self.path[-1], self.path[0]] if len(self.path) == self.num_nodes else float('inf')

class AntColonyOptimizer:
    def __init__(self, num_nodes, num_ants, num_iterations, alpha=1.0, beta=2.0, evaporation_rate=0.5, pheromone_deposit=1.0):
        self.num_nodes = num_nodes
        self.num_ants = num_ants
        self.num_iterations = num_iterations
        self.alpha = alpha
        self.beta = beta
        self.evaporation_rate = evaporation_rate
        self.pheromone_deposit = pheromone_deposit
        self.pheromones = np.ones((num_nodes, num_nodes))
        self.best_path = None
        self.best_cost = float('inf')

    def optimize(self, cost_matrix):
        ants = [Ant(self.num_nodes) for _ in range(self.num_ants)]
        for iteration in range(self.num_iterations):
            for ant in ants:
                ant.reset()
                start_node = np.random.randint(self.num_nodes)
                ant.visit_node(start_node, cost_matrix)
                while len(ant.path) < self.num_nodes:
                    current_node = ant.path[-1]
                    probabilities = self._calculate_transition_probabilities(ant, current_node, cost_matrix)
                    next_node = np.random.choice(range(self.num_nodes), p=probabilities)
                    ant.visit_node(next_node, cost_matrix)
                cost = ant.path_cost(cost_matrix)
                if cost < self.best_cost:
                    self.best_cost = cost
                    self.best_path = ant.path.copy()
            self._update_pheromones(ants, cost_matrix)
        return self.best_path, self.best_cost

    def _calculate_transition_probabilities(self, ant, current_node, cost_matrix):
        """
        路径选择规则
        """

        pheromones = self.pheromones[current_node]
        inverse_costs = 1.0 / cost_matrix[current_node]
        probabilities = (pheromones ** self.alpha) * (inverse_costs ** self.beta)
        probabilities[ant.visited] = 0
        total = np.sum(probabilities)
        if total == 0:
            return probabilities
        return probabilities / total

    def _update_pheromones(self, ants, cost_matrix):
        """
        信息素更新规则
        """
        self.pheromones *= (1 - self.evaporation_rate)
        for ant in ants:
            for i in range(len(ant.path) - 1):
                self.pheromones[ant.path[i], ant.path[i+1]] += self.pheromone_deposit / ant.path_cost(cost_matrix)#蚁周
            self.pheromones[ant.path[-1], ant.path[0]] += self.pheromone_deposit / ant.path_cost(cost_matrix)

def calculate_distance_matrix(coordinates):
    num_nodes = len(coordinates)
    distance_matrix = np.zeros((num_nodes, num_nodes))
    for i in range(num_nodes):
        for j in range(i + 1, num_nodes):
            distance = np.linalg.norm(coordinates[i] - coordinates[j])
            distance_matrix[i, j] = distance_matrix[j, i] = distance
    return distance_matrix

def read_coordinates_from_csv(file_path):
    df = pd.read_csv(file_path)
    coordinates = df[['x', 'y']].values
    return coordinates

def plot_path(coordinates, path):
    plt.figure()
    for i in range(len(path)):
        start = coordinates[path[i]]
        end = coordinates[path[(i + 1) % len(path)]]
        plt.plot([start[0], end[0]], [start[1], end[1]], 'bo-')
    plt.show()

if __name__ == "__main__":
    file_path = 'coordinates.csv'  # 修改为你的CSV文件路径
    coordinates = read_coordinates_from_csv(file_path)
    distance_matrix = calculate_distance_matrix(coordinates)

    num_ants = 10
    num_iterations = 100
    aco = AntColonyOptimizer(num_nodes=len(coordinates), num_ants=num_ants, num_iterations=num_iterations)
    best_path, best_cost = aco.optimize(distance_matrix)

    print("Best Path:", best_path)
    print("Best Cost:", best_cost)

    plot_path(coordinates, best_path)
